import json
import os
import random
import sys
from datetime import datetime

import httpx
from dotenv import load_dotenv
from openai import OpenAI

from mcp_demo.common.HttpLogger import HttpLogger

load_dotenv()
# 初始化日志记录器
HttpLogger.initialize(__file__, "./log")
http_client = httpx.Client(event_hooks={
    'request': [HttpLogger.log_request],
    'response': [HttpLogger.log_response]}
)

client = OpenAI(
    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
    api_key=os.getenv("MCP_DEMO_API_KEY"),
    base_url=os.getenv("MCP_DEMO_BASE_URL"),  # 填写DashScope SDK的base_url,
    http_client=http_client,
    default_headers={"Content-Type": "application/json; charset=utf-8"}
)

# 定义工具列表，模型在选择使用哪个工具时会参考工具的name和description
tools = [
    # 工具1 获取当前时刻的时间
    {
        "type": "function",
        "function": {
            "name": "get_current_time",
            "description": "当你想知道现在的时间时非常有用。",
            # 因为获取当前时间无需输入参数，因此parameters为空字典
            "parameters": {}
        }
    },
    # 工具2 获取指定城市的天气
    {
        "type": "function",
        "function": {
            "name": "get_current_weather",
            "description": "当你想查询指定城市的天气时非常有用。",
            "parameters": {
                "type": "object",
                "properties": {
                    # 查询天气时需要提供位置，因此参数设置为location
                    "location": {
                        "type": "string",
                        "description": "城市或县区，比如北京市、杭州市、余杭区等。"
                    }
                },
                "required": [
                    "location"
                ]
            }
        }
    }
]


# 模拟天气查询工具。返回结果示例：“北京今天是雨天。”
def get_current_weather(arguments):
    # 定义备选的天气条件列表
    weather_conditions = ["晴天", "多云", "雨天"]
    # 随机选择一个天气条件
    random_weather = random.choice(weather_conditions)
    # 从 JSON 中提取位置信息
    location = arguments["location"]
    # 返回格式化的天气信息
    return f"{location}今天是{random_weather}。"


# 查询当前时间的工具。返回结果示例：“当前时间：2024-04-15 17:15:18。“
def get_current_time():
    # 获取当前日期和时间
    current_datetime = datetime.now()
    # 格式化当前日期和时间
    formatted_time = current_datetime.strftime('%Y-%m-%d %H:%M:%S')
    # 返回格式化后的当前时间
    return f"当前时间：{formatted_time}。"


# 封装模型响应函数
def get_response(messages, parallel):
    print(f"请求messages {messages}")
    completion = client.chat.completions.create(
        model=os.getenv("MCP_DEMO_MODEL_NAME"),  # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
        messages=messages,
        tools=tools,
        parallel_tool_calls=parallel
    )
    return completion


def call_with_messages(parallel):
    print("并行工具调用" if parallel else "非并行工具调用")
    messages = [
        {
            "content": input('请输入：'),  # 提问示例："现在几点了？" "一个小时后几点" "北京天气如何？"
            "role": "user"
        }
    ]
    print("-" * 60)
    # 模型的第一轮调用
    i = 1
    response = get_response(messages, parallel)
    assistant_output = response.choices[0].message
    print(f"\n第{i}轮大模型输出信息：{response}\n")
    if assistant_output.content is None:
        assistant_output.content = ""
    # 如果不需要调用工具，则直接返回最终答案
    if response.choices[0].finish_reason != "tool_calls" \
            or assistant_output.tool_calls is None \
            or len(assistant_output.tool_calls) == 0:
        # 如果模型判断无需调用工具，则将assistant的回复直接打印出来，无需继续调用模型
        print(f"无需调用工具，我可以直接回复：{assistant_output.content}")
        return
    # 请求messages增加上一次大模型返回的choices[0].message，即role=assistant的内容
    messages.append(assistant_output)
    # 如果需要调用工具，则进行模型的多轮调用，直到模型判断无需调用工具
    while response.choices[0].finish_reason == "tool_calls" \
            and assistant_output.tool_calls is not None \
            and len(assistant_output.tool_calls) > 0:
        for tool_call in assistant_output.tool_calls:
            # 遍历所有工具调用
            tool_name = tool_call.function.name
            arguments = json.loads(tool_call.function.arguments)

            # 执行对应工具
            if tool_name == "get_current_weather":
                result = get_current_weather(arguments)
            elif tool_name == "get_current_time":
                result = get_current_time()
            else:
                result = "未知工具"

            # 处理调用工具的结果，tool_call_id 字段设置为当前处理的 tool_calls 数组元素的 id
            messages.append({
                "role": "tool",
                "content": result,
                "tool_call_id": tool_call.id
            })
            if not parallel:
                break

        response = get_response(messages, parallel)
        assistant_output = response.choices[0].message
        if assistant_output.content is None:
            assistant_output.content = ""
        messages.append(assistant_output)
        i += 1
        print(f"第{i}轮大模型输出信息：{assistant_output}\n")
    print(f"最终答案：{assistant_output.content}")


if __name__ == '__main__':
    # 只要参数1非空，bool()就会返回True，代表并行工具调用
    parallel = bool(sys.argv[1]) if len(sys.argv) > 1 else False
    call_with_messages(parallel)
